import pinocchio as pin
from pinocchio.utils import rand, zero
import meshcat

# 引入src中所有的函数
import sys
from os.path import dirname, join

sys.path.append(join(dirname(dirname(__file__)), "src"))
sys.path.append(dirname(dirname(__file__))) # 把文件根路径添加到sys.path中，这样在根路径下的src文件夹才能被import
from src import *
# <<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<

# 获取URDF模型路径
URDF_FILE_PATH = filePath(URDF_NAME = 'ur5_gripper', ROBOT = True)
# <<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<

# 构建robot对象，并得到viz
robot, viz = robot_viz(URDF_FILE_PATH)
q0 = pin.neutral(robot.model)
pin.forwardKinematics(robot.model, robot.data, q0)
robot.initViewer(loadModel=True)
viz.display(q0)
# <<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<
# <<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<



# <<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<
# 开始进行逆运动学求解
# 此程序只解决位置，不解决姿态，位姿IK请查看inverse_kinematics.py
# 导入求解优化问题必要的库
import numpy as np
from scipy.optimize import fmin_bfgs


# 设置effector目标位置
destination = np.array([0.7, 0.0, 0.3])
############################################################
# 创建一个box位于目标位置，方便查看
box_name = "world/box"
box_size = [0.04, 0.02, 0.01]

viz.viewer[box_name].set_object(
    meshcat.geometry.Box(box_size),
    meshcat.geometry.MeshLambertMaterial(color=0x6622cc,reflectivity=0.4)
)

# 设置box的位置和姿态
position = destination
quaternion = [0.5, 0.5, 0.5, 0.5]  # w, x, y, z
rotation_matrix = pin.Quaternion(quaternion[0],quaternion[1],quaternion[2],quaternion[3]).matrix()
viz.viewer[box_name].set_transform(
    pin.SE3(rotation_matrix, position).homogeneous
)

# <<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<




# 优化末端执行器的位置（三维坐标），逆运动学
# 定义损失函数
def cost(q):
    pin.forwardKinematics(robot.model, robot.data, q)
    p = robot.data.oMi[6].translation
    return np.linalg.norm(p - destination)

xopt_bfgs = fmin_bfgs(cost, q0)
print('*** Xopt in BFGS =', xopt_bfgs)
viz.display(xopt_bfgs)
# <<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<


# # 优化末端执行器的姿态(六维坐标),SE3，逆运动学(优化失败，逆运动学不应当直接选择bfgs算法)
# # 此处选择 log in SE3来定义SE3群中的度量（metric）
# from pinocchio.explog import log
# from pinocchio import SE3

# destination_se3 = log(pin.SE3(rotation_matrix, position)).vector

# def cost_se3(q):
#     pin.forwardKinematics(robot.model, robot.data, q)
#     p_se3 = log(robot.data.oMi[6]).vector
#     # return p_se3 - destination_se3
#     return np.linalg.norm(p_se3 - destination_se3)

# xopt_bfgs_se3 = fmin_bfgs(cost_se3, q0)
# print('*** Xopt in BFGS =', xopt_bfgs_se3)
# viz.display(xopt_bfgs_se3)
# # <<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<